package eval;

import inference.TableInferer;

import java.io.BufferedReader;
import java.io.BufferedWriter;
import java.io.FileNotFoundException;
import java.io.FileReader;
import java.io.FileWriter;
import java.io.IOException;
import java.util.ArrayList;
import java.util.Collections;
import java.util.Comparator;
import java.util.HashMap;
import java.util.List;

import model.TableInformation;
import utils.Pair;

/**
 * InferenceEvaluator runs TableInferer on a table or list of tables and writes to file information
 * about each run.
 * 
 * Command line arguments are written in Flag format.
 * Arguments:
 * 		--table=table_name
 * 		--tables_file=file_of_table_names
 * 		--data_dir=/path/to/input/files/
 * 		--output_dir=/path/to/output/files/
 * 		--extra_info=information_about_the_run_to_be_put_in_output_filenames
 * 		--max_hidden_cols=number of allowable hidden CVT columns
 * 
 * Outputs:
 * 		For a table named Table,
 * 			Table.extra_info.output.info : contains information about the speed of the run, the backbones checked,
 * 				and the probabilities for predicates and entities in the best backbone
 * 			Table.extra_info.output.facts : a list of the best triples inferred for the table using the best predicate
 * 				and entity choices from the best backbone's interpretation.
 * @author chloe
 *
 */
public class InfererenceEvaluator {
	private String tables_filename = "";
	private String table = "";
	private String data_dir = "";
	private String output_dir = "";
	private String extra_info = "";
	private int max_hidden_cols = 0;
	
	private List<String> table_filenames;
	
	private static String relations_to_mediators_filename = "/Users/chloe/Research/tablejoin/full_data/relations_to_mediators.txt";
	private HashMap<String, String> flags = new HashMap<String, String>();
	
	private static String getEntityToTypesFilename(String table_name) {
		return table_name + ".all.types";
	}
	
	private static String getFactFilename(String table_name) {
		return table_name + ".facts";
	}
	
	private static String getPropertyInfoFilename(String table_name) {
		return table_name + ".property_info";
	}
	
	private String getExperimentOutputFilename(String table_name, String info_name) {
		if (info_name.equals("")) {
			return output_dir + table_name + ".output.info";
		} else {
			return output_dir + table_name + "." + info_name + ".output.info";
		}
	}
	
	private String getExperimentOutputFactsFilename(String table_name, String info_name) {
		if (info_name.equals("")) {
			return output_dir + table_name + ".output.facts";
		} else {
			return output_dir + table_name + "." + info_name + ".output.facts";
		}
	}
	
	private Pair<Double, Double> runInferenceForTable(TableInferer inferer, String table_name, String entity_filename, String type_filename,
			String fact_filename, String predicate_info_filename, String predicate_to_mediator_filename) throws IOException {
		long start_setup = System.nanoTime();
		inferer.setup(entity_filename, type_filename, fact_filename, predicate_info_filename, predicate_to_mediator_filename, "");
		long end_setup = System.nanoTime();
		long start_inf = System.nanoTime();
		inferer.run();
		long end_inf = System.nanoTime();
		double setup_time = (end_setup - start_setup) / ((double)1000000000);
		double inference_time = (end_inf - start_inf)/ ((double)1000000000);
		return new Pair<Double, Double>(setup_time, inference_time);
	}
	
	private boolean runInference() {
		for (String table : table_filenames) {
			try {
				TableInferer inferer = new TableInferer();
				String entity_filename = data_dir + table;
				System.out.println(entity_filename);
				String type_filename = data_dir + getEntityToTypesFilename(table);
				String fact_filename = data_dir + getFactFilename(table);
				String predicate_info_filename = data_dir + getPropertyInfoFilename(table);
				Pair<Double, Double> times = runInferenceForTable(inferer, table, entity_filename, type_filename, fact_filename, predicate_info_filename, relations_to_mediators_filename);
				writeInferenceOutputForTable(inferer, table, times);
			} catch (IOException ex) {
				String output_filename = getExperimentOutputFilename(table, extra_info);
				try {
					BufferedWriter writer = new BufferedWriter(new FileWriter(output_filename));
					writer.write("No interpretation.\n");
					writer.close();
				} catch (IOException ex2) {
					ex2.printStackTrace();
				}
				ex.printStackTrace();
				return false;
			}
		}
		return true;
	}
	
	private String expandIfUnkString(String entity, String cell_text, int row, int col) {
		if (entity.equals("UNKENT")) {
			return entity + "[" + row + "," + col + "](" + cell_text + ")"; 
		} else {
			return entity;
		}
	}
	
	private ArrayList<Pair<String, Double> > selectBestFacts(int[] backbone, int backbone_length, int num_rows, TableInformation table_info,
			String best_topic, double best_topic_prob,
			String[] best_preds, double[] best_pred_probs, String[][] best_entities, double[][] best_entity_probs) {
		ArrayList<Pair<String, Double> > facts = new ArrayList<Pair<String, Double> >();
		String[][] cell_text = table_info.getCellText();
		
		for (int col = 0; col < backbone_length; col++) {
			int par_col = backbone[col];
			String best_pred = best_preds[col];
			double best_pred_prob = best_pred_probs[col];
			if (par_col == -1) {
				String par_entity = best_topic;
				for (int row = 0; row < num_rows; row++) {
					String entity;
					if (col >= table_info.numVisibleCols()) {
						entity = expandIfUnkString(best_entities[row][col], "", row, col);
					} else {
						entity = expandIfUnkString(best_entities[row][col], cell_text[col][row], row, col);
					}
					boolean is_known_fact = table_info.isKnownFact(par_entity, best_pred, entity);
					String fact = par_entity + " " + best_pred + " " + entity;
					if (is_known_fact) {
						fact = "*" + fact;
					}
					facts.add(new Pair<String, Double>(fact,
							(best_topic_prob * best_pred_prob * best_entity_probs[row][col])));
				}
			} else {
				for (int row = 0; row < num_rows; row++) {
					String par_entity;
					if (par_col >= table_info.numVisibleCols()) {
						System.out.println(par_col);
						par_entity = expandIfUnkString(best_entities[row][par_col], "", row, par_col);
					} else {
						par_entity = expandIfUnkString(best_entities[row][par_col], cell_text[par_col][row], row, par_col);
					}
					String entity;
					if (col >= table_info.numVisibleCols()) {
						entity = expandIfUnkString(best_entities[row][col], "", row, col);
					} else {
						entity = expandIfUnkString(best_entities[row][col], cell_text[col][row], row, col);
					}
					boolean is_known_fact = table_info.isKnownFact(par_entity, best_pred, entity);
					String fact = par_entity + " " + best_pred + " " + entity;
					if (is_known_fact) {
						fact = "*" + fact;
					}
					facts.add(new Pair<String, Double>(fact,
							(best_entity_probs[row][par_col] * best_pred_prob * best_entity_probs[row][col])));
				}
			}
		}
		Collections.sort(facts, new Comparator<Pair<String, Double> >()
        {
            public int compare(Pair<String, Double>o1,
            		Pair<String, Double> o2)
            {
                    return o2.getSecond().compareTo(o1.getSecond());
            }
        });
		return facts;
	}
	
	private void writeInferenceOutputForTable(TableInferer inferer, String table_name, Pair<Double, Double> times) {
		String output_filename = getExperimentOutputFilename(table_name, extra_info);
		String fact_filename = getExperimentOutputFactsFilename(table_name, extra_info);
		TableInformation table_info = inferer.getModelParameters().getTableInformation();
		if (times == null) {
			try {
				BufferedWriter writer = new BufferedWriter(new FileWriter(output_filename));
				writer.write("No interpretation.\n");
				writer.close();
			} catch (IOException ex) {
				ex.printStackTrace();
			}
		} else {
		try {
			BufferedWriter writer = new BufferedWriter(new FileWriter(output_filename));
			
			// Write down the best bethe free energy.
			writer.write(inferer.getBestBetheFreeEnergy() + "\n");
			
			// Write down the timing.
			writer.write(times.getFirst().toString() + " " + times.getSecond().toString() + "\n");
			
			// Write down all backbones checked.
			ArrayList<Pair<int[], Double> > backbones_checked = inferer.getBackbonesChecked();
			for (Pair<int[], Double> checked_backbone : backbones_checked) {
				int[] backbone = checked_backbone.getFirst();
				double score = checked_backbone.getSecond();
				
				writer.write(score + "\t\t");
				for (int c = 0; c < backbone.length; c++) {
					writer.write(backbone[c] + "");
					if (c == backbone.length - 1) {
						writer.write("\n");
					} else {
						writer.write(" ");
					}
				}
			}
			
			// Write down the backbone.
			int[] backbone = inferer.getBestBackbone();
			int backbone_length = inferer.getBestBackboneLength();
			System.out.println(backbone_length);
			for (int i = 0; i < backbone_length; i++) {
				writer.write(backbone[i] + "");
				if (i == backbone_length - 1) {
					writer.write("\n");
				} else {
					writer.write(" ");
				}
			}
			
			// Write sorted topics.
			Pair<String[], double[]> sorted_topics = inferer.getSortedTopics();
			String[] sorted_topic_domain = sorted_topics.getFirst();
			double[] sorted_topic_probs = sorted_topics.getSecond();
			for (int i = 0; i < sorted_topic_domain.length; i++) {
				writer.write(sorted_topic_domain[i] + ":" + sorted_topic_probs[i]);
				if (i == sorted_topic_domain.length - 1) {
					writer.write("\n");
				} else {
					writer.write(" ");
				}
			}
			writer.write("\n");
			
			String[] best_preds = new String[backbone_length];
			double[] best_pred_probs = new double[backbone_length];
			// Write sorted column predicates.
			for (int col = 0; col < backbone_length; col++) {
				writer.write(col + " ");
				System.out.println(col);
				Pair<String[], double[]> sorted_preds = inferer.getSortedPredicatesForColumn(col);
				String[] sorted_pred_domain = sorted_preds.getFirst();
				double[] sorted_pred_probs = sorted_preds.getSecond();
				for (int i = 0; i < sorted_pred_domain.length; i++) {
					writer.write(sorted_pred_domain[i] + ":" + sorted_pred_probs[i]);
					if (i == sorted_pred_domain.length - 1) {
						writer.write("\n");
					} else {
						writer.write(" ");
					}
				}
				best_preds[col] = sorted_pred_domain[0];
				best_pred_probs[col] = sorted_pred_probs[0];
			}
			
			int num_rows = table_info.numRows();
			String[][] best_entities = new String[num_rows][backbone_length];
			double[][] best_entity_probs = new double[num_rows][backbone_length];
			for (int row = 0; row < num_rows; row++) {
				for (int col = 0; col < backbone_length; col++) {
					Pair<String[], double[]> sorted_entities = inferer.getSortedEntitiesForCell(row, col);
					String[] sorted_entity_domain = sorted_entities.getFirst();
					double[] sorted_entity_probs = sorted_entities.getSecond();
					writer.write("(" + row + "," + col + ") ");
					for (int i = 0; i < sorted_entity_domain.length; i++) {
						writer.write(sorted_entity_domain[i] + ":" + sorted_entity_probs[i]);
						if (i == sorted_entity_domain.length - 1) {
							writer.write("\n");
						} else {
							writer.write(" ");
						}
					}
					best_entities[row][col] = sorted_entity_domain[0];
					best_entity_probs[row][col] = sorted_entity_probs[0];
				}
			}
			
			writer.close();
			ArrayList<Pair<String, Double> > best_facts = 
					selectBestFacts(backbone, backbone_length, num_rows, table_info, sorted_topic_domain[0], sorted_topic_probs[0], best_preds, best_pred_probs,
							best_entities, best_entity_probs);
			writer = new BufferedWriter(new FileWriter(fact_filename));
			for (Pair<String, Double> fact : best_facts) {
				writer.write(fact.getSecond() + " " + fact.getFirst() + "\n");
			}
			writer.close();
		} catch (IOException e) {
			e.printStackTrace();
		}
		}
	}

	private boolean readInTableNames() {
		if (tables_filename.equals("") && table.equals("")) {
			System.err.println("Please run readInCommandLineFlags(String[] args) before calling readITableNames().");
			return false;
		}
		table_filenames = new ArrayList<String>();
		if (!table.equals("")) {
			table_filenames.add(table);
		} else {
			try {
				BufferedReader reader = new BufferedReader(new FileReader(tables_filename));
				String line;
				while ((line = reader.readLine()) != null) {
					table_filenames.add(line);
				}
				reader.close();
			} catch (FileNotFoundException e) {
				e.printStackTrace();
				return false;
			} catch (IOException e) {
				// TODO Auto-generated catch block
				e.printStackTrace();
				return false;
			}
		}
		
		return true;
	}
	
	private boolean readInCommandLineFlags(String[] args) {
		for (int a = 0; a < args.length; a++) {
			String arg = args[a];
			String[] split = arg.split("=");
			if (split.length == 2) {
				if (split[0].equals("--tables_file")) {
					if (split[1].equals("")) {
						System.err.println("Missing flag argument: --tables_file");
						return false;
					}
					tables_filename = split[1];
				} else if (split[0].equals("--table")) {
					if (split[1].equals("")) {
						System.err.println("Missing flag argument: --table");
						return false;
					}
					table = split[1];
				} else if (split[0].equals("--data_dir")) {
					if (split[1].equals("")) {
						System.err.println("Missing flag argument: --data_dir");
						return false;
					}
					data_dir = split[1];
				} else if (split[0].equals("--output_dir")) {
					if (split[1].equals("")) {
						System.err.println("Missing flag argument: --output_dir");
						return false;
					}
					output_dir = split[1];
				} else if (split[0].equals("--extra_info")) {
					if (split[1].equals("")) {
						System.err.println("Missing flag argument: --extra_info");
						return false;
					}
					extra_info = split[1];
				} else if (split[0].equals("--relations_to_mediators")) {
					if (split[1].equals("")) {
						System.err.println("Missing flag argument: --relations_to_mediators");
						return false;
					}
					relations_to_mediators_filename = split[1];
				} else if (split[0].equals("--max_hidden_cols")) {
					if (split[1].equals("")) {
						System.err.println("Missing flag argument: --max_hidden_cols");
						return false;
					}
					max_hidden_cols = Integer.parseInt(split[1]);
				} else {
					if (!split[0].startsWith("--")) {
						System.err.println("Invalid flag: " + arg);
						return false;
					}
					if (split[1].equals("")) {
						System.err.println("Missing flag argument: " + arg);
						return false;
					}
					flags.put(split[0], split[1]);
				}
			}
		}
		System.out.println(table_filenames);
		System.out.println(table + "..");
		if (tables_filename.equals("") && table.equals("")) {
			System.err.println("Please specify either a file of tables (--tables_filename) or a table file (--table).");
			return false;
		}
		if (!tables_filename.equals("") && !table.equals("")) {
			System.err.println("Please specify either a file of tables (--tables_filename) or a table file (--table).");
			return false;
		}
		if (!data_dir.equals("") && !data_dir.endsWith("/")) {
			data_dir = data_dir + "/";
		}
		if (!output_dir.equals("") && !output_dir.endsWith("/")) {
			output_dir = output_dir + "/";
		}
		return true;
	}
	
	public static void main(String[] args) {
		InfererenceEvaluator evaluator = new InfererenceEvaluator();
		boolean read_args_correctly = evaluator.readInCommandLineFlags(args);
		if (!read_args_correctly) {
			return;
		}
		boolean read_table_names_correctly = evaluator.readInTableNames();
		if (!read_table_names_correctly) {
			return;
		}
		evaluator.runInference();
	}
}
